#include "tools.h"

// 实现简单的汉宁窗函数
std::vector<float> hanningWindow(int size) {
	std::vector<float> window(size);
	for (int i = 0; i < size; ++i) {
		window[i] = 0.5f * (1.f - std::cos(2.f * M_PI * i / (size - 1)));
	}
	return window;
}

// display progress bar
void displayProgressBar(int progress, int total) {
	const int barWidth = 70;

	float fraction = static_cast<float>(progress) / total;
	int barLength = static_cast<int>(fraction * barWidth);

	std::cout << "[";

	for (int i = 0; i < barLength; ++i) {
		std::cout << "*";
	}
	for (int i = barLength; i < barWidth; ++i) {
		std::cout << " ";
	}

	std::cout << "] " << int(fraction * 100.0) << "%\r";
	std::cout.flush();
}

// determine whether file exist
bool fileExists(const std::string& filename) {
	std::ifstream file(filename);
	return file.good();
}

// 实现简化的STFT函数
std::vector<std::vector<Complex>> simplifiedSTFT(const std::vector<float>& x, int n_fft, int hop_length, const std::string& window, bool center, const std::string& mode) {
	std::vector<float> win;
	if (window == "hann") {
		win = hanningWindow(n_fft);
	}
	else {
		// 如果需要，添加其他类型的窗函数
	}

	// 如果需要，进行填充
	int pad_len = center ? n_fft / 2 : 0;
	std::vector<float> x_padded(x.size() + 2 * pad_len, 0);
	std::copy(x.begin(), x.end(), x_padded.begin() + pad_len);

	int n_frames = 1 + (x_padded.size() - n_fft) / hop_length;
	std::vector<std::vector<Complex>> X(n_frames, std::vector<Complex>(n_fft / 2 + 1, 0));

	for (int i = 0; i < n_frames; ++i) {
		std::vector<float> x_frame(n_fft);
		for (int j = 0; j < n_fft; ++j) {
			x_frame[j] = win[j] * x_padded[i * hop_length + j];
		}

		// 执行快速傅立叶变换
		for (int k = 0; k < n_fft / 2 + 1; ++k) {
			Complex sum = 0;
			for (int j = 0; j < n_fft; ++j) {
				sum += Complex(x_frame[j] * std::cos(2 * M_PI * j * k / n_fft),
					-x_frame[j] * std::sin(2 * M_PI * j * k / n_fft));
			}
			X[i][k] = sum;
		}
	}

	return X;
}

std::vector<std::vector<float>> melFilterBank(int n_fft, int n_mel, int sample_rate)
{
	// 计算频率对应的Mel频率
	std::vector<float> mel_frequencies(n_mel + 2);
	float mel_min = 0;
	float mel_max = 2595 * log10(1 + sample_rate / 2.f / 700);
	for (int i = 0; i < n_mel + 2; ++i) {
		mel_frequencies[i] = mel_min + i * (mel_max - mel_min) / (n_mel + 1);
	}

	// 将Mel频率转换为线性频率
	std::vector<float> frequencies(n_mel + 2);
	for (int i = 0; i < n_mel + 2; ++i) {
		frequencies[i] = 700 * (pow(10, mel_frequencies[i] / 2595) - 1);
	}

	// 初始化梅尔滤波器组
	std::vector<std::vector<float>> filters(n_mel, std::vector<float>(n_fft / 2 + 1, 0));

	for (int i = 0; i < n_mel; ++i) {
		float mel_center = mel_frequencies[i + 1];
		float mel_left = mel_frequencies[i];
		float mel_right = mel_frequencies[i + 2];
		float f_center = frequencies[i + 1];
		float f_left = frequencies[i];
		float f_right = frequencies[i + 2];

		for (int j = 0; j <= n_fft / 2; ++j) {
			float freq = sample_rate * j / static_cast<float>(n_fft);
			if (freq >= f_left && freq <= f_right) {
				if (freq <= f_center)
					filters[i][j] = (freq - f_left) / (f_center - f_left);
				else
					filters[i][j] = (f_right - freq) / (f_right - f_center);
			}
		}
	}

	return filters;
}

std::vector<std::vector<float>> melSpectrogram(const std::vector<float>& x, int sample_rate, int n_fft, int n_hop, const std::string& win, bool center, const std::string& mode, float ref_power, int n_mel, int fmin, int fmax)
{
	// 计算STFT
	std::vector<float> window = hanningWindow(n_fft);
	int pad_len = center ? n_fft / 2 : 0;
	std::vector<float> x_padded(x.size() + 2 * pad_len, 0);
	std::copy(x.begin(), x.end(), x_padded.begin() + pad_len);
	int n_frames = 1 + (x_padded.size() - n_fft) / n_hop;
	std::vector<std::vector<Complex>> stft_result(n_frames, std::vector<Complex>(n_fft / 2 + 1, 0));
	for (int i = 0; i < n_frames; ++i) {
		std::vector<float> x_frame(n_fft);
		for (int j = 0; j < n_fft; ++j) {
			x_frame[j] = window[j] * x_padded[i * n_hop + j];
		}
		for (int k = 0; k < n_fft / 2 + 1; ++k) {
			Complex sum = 0;
			for (int j = 0; j < n_fft; ++j) {
				sum += Complex(x_frame[j] * std::cos(2 * M_PI * j * k / n_fft),
					-x_frame[j] * std::sin(2 * M_PI * j * k / n_fft));
			}
			stft_result[i][k] = sum;
		}
	}

	// 计算梅尔滤波器组
	std::vector<std::vector<float>> filters = melFilterBank(n_fft, n_mel, sample_rate);

	// 只选择感兴趣的频率范围内的数据
	int n_fmin_bin = static_cast<int>(std::floor(fmin * n_fft / sample_rate));
	int n_fmax_bin = static_cast<int>(std::ceil(fmax * n_fft / sample_rate)) + 1;
	std::vector<std::vector<Complex>> stft_result_selected(n_frames, std::vector<Complex>(n_fmax_bin - n_fmin_bin, 0));
	for (int i = 0; i < n_frames; ++i) {
		std::copy(stft_result[i].begin() + n_fmin_bin, stft_result[i].begin() + n_fmax_bin, stft_result_selected[i].begin());
	}

	// 计算梅尔频谱图
	std::vector<std::vector<float>> mel_spectrogram(n_frames, std::vector<float>(n_mel, 0));
	for (int i = 0; i < n_frames; ++i) {
		for (int j = 0; j < n_mel; ++j) {
			float sum = 0;
			for (int k = 0; k < stft_result[i].size(); ++k) {
				sum += std::abs(stft_result[i][k]) * filters[j][k];
			}
			mel_spectrogram[i][j] = sum;
		}
	}

	return mel_spectrogram;
}
